# Testing for conformation embedding.
# 23-11-01

import os
import sys
import jax
import jax.numpy as jnp
import numpy as np
import haiku as hk
import mindspore as ms
import mindspore.nn as ms_nn

from jax import Array, jit
from typing import Optional, Union, List
from jax.nn.initializers import lecun_normal, normal, zeros, ones
from mindspore import load_param_into_net, Tensor, nn, context

sys.path.append(os.path.dirname(sys.path[0]))

# setting context & hyper parameters
context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU")

from cybertron.embedding import ConformationEmbedding as Emb
from ms_cybertron.embedding import MolEmbedding as ms_Emb

EPSILON = 1e-3
rng = jax.random.PRNGKey(42)

emb_data = np.load(sys.path[0] + '/conformation_data.npz')
atom_type = emb_data['atom_type']
atom_mask = emb_data['atom_mask']
bond_type = None
bond_mask = emb_data['bond_mask']
dist = emb_data['dist'] / 10.0
dist_mask = emb_data['dist_mask']

print("=================Test ConformationEmbedding====================")

def emb_fn(a, b, c, d, e, f):
    emb = Emb(dim_node=16, dim_edge=16, num_atom_types=64, rbf_runc='gaussian', dis_self=0.05)
    return emb(a, b, c, d, e, f)

emb_fn = jax.vmap(emb_fn, in_axes=(0, 0, 0, 0, 0, 0))
ts_emb_fn = hk.transform(emb_fn, apply_rng=True)
emb_params = ts_emb_fn.init(rng, atom_type, bond_type, dist, atom_mask, bond_mask, dist_mask)
node_emb, node_mask, edge_emb, edge_mask, edge_cutoff = jit(ts_emb_fn.apply)(emb_params, rng,
                atom_type, bond_type, dist, atom_mask, bond_mask, dist_mask)
# print("[from jax]")
# print(f"node_emb with shape of {node_emb.shape}")
# print(f"node_mask with shape of {node_mask.shape}")
# print(f"edge_emb with shape of {edge_emb.shape}")
# print(f"edge_mask with shape of {edge_mask.shape}")
# print(f"edge_cutoff with shape of {edge_cutoff.shape}")

ms_emb_fn = ms_Emb(dim_node=16, dim_edge=16, num_atom_types=64, dis_filter='residual', rbf_fn='gaussian')
ms_params = {}
_params = hk.data_structures.to_mutable_dict(emb_params)
ms_params['atom_embedding.embedding_table'] = \
    ms.Parameter(Tensor(np.array(_params[f'conformation_embedding/~/atom_embedding']['embeddings']), ms.float32))
ms_params['dis_filter.linear.weight'] = \
    ms.Parameter(Tensor(np.array(_params[f'conformation_embedding/~/dis_filter/~/linear']['w']).T, ms.float32))
ms_params['dis_filter.linear.bias'] = \
    ms.Parameter(Tensor(np.array(_params[f'conformation_embedding/~/dis_filter/~/linear']['b']), ms.float32))
ms_params['dis_filter.residual.nonlinear.mlp.0.weight'] = \
    ms.Parameter(Tensor(np.array(_params[f'conformation_embedding/~/dis_filter/~/mlp/~/linear_0']['w']).T, ms.float32))
ms_params['dis_filter.residual.nonlinear.mlp.0.bias'] = \
    ms.Parameter(Tensor(np.array(_params[f'conformation_embedding/~/dis_filter/~/mlp/~/linear_0']['b']), ms.float32))
ms_params['dis_filter.residual.nonlinear.mlp.1.weight'] = \
    ms.Parameter(Tensor(np.array(_params[f'conformation_embedding/~/dis_filter/~/mlp/~/linear_1']['w']).T, ms.float32))
ms_params['dis_filter.residual.nonlinear.mlp.1.bias'] = \
    ms.Parameter(Tensor(np.array(_params[f'conformation_embedding/~/dis_filter/~/mlp/~/linear_1']['b']), ms.float32))
load_param_into_net(ms_emb_fn, ms_params)

ms_atom_type = Tensor(atom_type, ms.int32)
ms_atom_mask = Tensor(atom_mask, ms.bool_)
ms_bond_type = None
ms_bond_mask = Tensor(bond_mask, ms.bool_)
ms_dist = Tensor(dist, ms.float32)
ms_dist_mask = Tensor(dist_mask, ms.bool_)

ms_node_emb, ms_node_mask, ms_edge_emb, ms_edge_mask, ms_edge_cutoff = \
    ms_emb_fn(ms_atom_type, ms_dist, ms_bond_type, ms_atom_mask, ms_dist_mask, ms_bond_mask, )

# print("[from mindspore]")
# print(f"node_emb with shape of {ms_node_emb.shape}")
# print(f"node_mask with shape of {ms_node_mask.shape}")
# print(f"edge_emb with shape of {ms_edge_emb.shape}")
# print(f"edge_mask with shape of {ms_edge_mask.shape}")
# print(f"edge_cutoff with shape of {ms_edge_cutoff.shape}")

print("[compare]")
node_emb_same = np.allclose(node_emb, ms_node_emb.asnumpy(), atol=EPSILON)
node_emb_diff_norm = np.linalg.norm(node_emb - ms_node_emb.asnumpy())
print(f"node_emb_same: {node_emb_same} in accuracy of {EPSILON}, node_emb_diff_norm: {node_emb_diff_norm}")
node_mask_same = np.allclose(node_mask, ms_node_mask.asnumpy(), atol=EPSILON)
# node_mask_diff_norm = np.linalg.norm(node_mask - ms_node_mask.asnumpy())
print(f"node_mask_same: {node_mask_same} in accuracy of {EPSILON}")
edge_emb_same = np.allclose(edge_emb, ms_edge_emb.asnumpy(), atol=EPSILON)
edge_emb_diff_norm = np.linalg.norm(edge_emb - ms_edge_emb.asnumpy())
print(f"edge_emb_same: {edge_emb_same} in accuracy of {EPSILON}, edge_emb_diff_norm: {edge_emb_diff_norm}")
edge_mask_same = np.allclose(edge_mask, ms_edge_mask.asnumpy(), atol=EPSILON)
# edge_mask_diff_norm = np.linalg.norm(edge_mask - ms_edge_mask.asnumpy())
print(f"edge_mask_same: {edge_mask_same} in accuracy of {EPSILON}")
edge_cutoff_same = np.allclose(edge_cutoff, ms_edge_cutoff.asnumpy(), atol=EPSILON)
edge_cutoff_diff_norm = np.linalg.norm(edge_cutoff - ms_edge_cutoff.asnumpy())
print(f"edge_cutoff_same: {edge_cutoff_same} in accuracy of {EPSILON}, edge_cutoff_diff_norm: {edge_cutoff_diff_norm}")

print("[from jax] edge emb:")
print(edge_emb[0, 0, 0, :])
print("[from mindspore] edge emb:")
print(ms_edge_emb.asnumpy()[0, 0, 0, :])    